import os
import numpy as np
from PIL import Image, ImageEnhance
import base64
import json
import glob
import concurrent.futures
from copy import deepcopy as dopy
import time

python_file = os.path.dirname(__file__)
results = os.path.join(python_file, 'results')
original = os.path.join(python_file, 'original')
positions = ['left_right', 'top_bottom', 'center', 'normal']
methods = ['normal', 'bright', 'compare']
dirs = {}
for dir_0 in positions:
    dirs[dir_0] = {}
    for dir_1 in methods:
        dirs[dir_0][dir_1] = os.path.join(results, dir_0, dir_1)
        if not os.path.exists(dirs[dir_0][dir_1]):
            os.makedirs(dirs[dir_0][dir_1])
result_all = os.path.join(results, 'all')
if not os.path.exists(result_all):
    os.makedirs(result_all)

lower = 0.5
upper = 1.5

# 随机改变亮度
def random_brightness(img):
    e = np.random.uniform(lower, upper)
    img = ImageEnhance.Brightness(img).enhance(e)
    return img

# 随机改变对比度
def random_contrast(img):
    e = np.random.uniform(lower, upper)
    img = ImageEnhance.Contrast(img).enhance(e)
    return img

def position_transformer(img_, setting_, command):
    img = dopy(img_)
    setting = dopy(setting_)
    if command == 'CE_':
        img, setting = position_transformer(img, setting, 'TB_')
        img, setting = position_transformer(img, setting, 'LR_')
        setting['imagePath'] = setting['imagePath'].replace('LR_TB_', 'CE_')
        return img, setting
    elif command == 'LR_':
        axis_0 = 'imageWidth'
        axis_1 = 0
        img = img.transpose(Image.FLIP_LEFT_RIGHT)
    elif command == 'TB_':
        axis_0 = 'imageHeight'
        axis_1 = 1
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
    else:
        return img, setting
    maximum = setting[axis_0]
    for everyshape in setting['shapes']:
        points = np.array(everyshape['points']).T
        points[axis_1] = maximum - points[axis_1]
        everyshape['points'] = points.T.tolist()
    setting['imagePath'] = command + setting['imagePath']
    return img, setting

def save(img, setting, dir):
    img_path = os.path.join(dir, setting['imagePath'])
    all_path = os.path.join(result_all, setting['imagePath'])
    img.save(img_path)
    img.save(all_path)
    with open(img_path, 'rb') as f:
        setting['imageData'] = base64.b64encode(f.read()).decode()
    string = json.dumps(setting)
    with open(img_path.replace('jpg', 'json'), 'w', encoding='utf-8') as f:
        f.write(string)
    with open(all_path.replace('jpg', 'json'), 'w', encoding='utf-8') as f:
        f.write(string)

def Enlarger(img, setting):
    br_img = random_brightness(img)
    br_setting = setting.copy()
    br_setting['imagePath'] = 'BR_' + br_setting['imagePath']
    co_img = random_contrast(img)
    co_setting = setting.copy()
    co_setting['imagePath'] = 'CO_' + co_setting['imagePath']
    for img_, setting_, color_method in zip([img, br_img, co_img], [setting, br_setting, co_setting], methods):
        for command, position in zip(['LR_', 'TB_', 'CE_', ''], positions):
            i, s = position_transformer(img_, setting_, command)
            save(i, s, dirs[position][color_method])

def main(img_path):
    img = Image.open(img_path)
    with open(img_path.replace('.jpg', '.json'), 'r', encoding='utf-8') as f:
        setting = json.load(f)
    Enlarger(img, setting)

if __name__ == '__main__':
    start = time.time()
    img_list = glob.glob(os.path.join(original, '*.jpg'))
    json_list = glob.glob(os.path.join(original, '*.json'))
    with concurrent.futures.ProcessPoolExecutor() as executor:
        for img_path in img_list:
            if img_path.replace('.jpg', '.json') in json_list:
                executor.submit(main, img_path)
    end = time.time()
    print('TIME: ' + str(end - start))
    
